import base64
import requests
from app.core.config import settings
import asyncio
from typing import List, Dict, Any, Optional, Tuple
import os
from PIL import Image
import pytesseract
import traceback
import logging
import aiohttp
import json
import io
import random
import re
import yaml

class OCRClient:
    """OCR客户端，支持多种OCR方式"""
    
    @staticmethod
    def get_access_token() -> str:
        """
        获取百度OCR的access_token
        """
        url = "https://aip.baidubce.com/oauth/2.0/token"
        params = {
            "grant_type": "client_credentials", 
            "client_id": settings.BAIDU_API_KEY, 
            "client_secret": settings.BAIDU_SECRET_KEY
        }
        response = requests.post(url, params=params)
        return str(response.json().get("access_token"))
    
    @staticmethod
    async def recognize_image(image_path: str) -> Dict:
        """
        使用百度OCR API识别图片
        
        Args:
            image_path: 图片路径
            
        Returns:
            Dict: OCR识别结果
        """
        try:
            from aiohttp import ClientSession
            from app.core.config import settings
            import base64
            import time
            
            # 读取图片文件
            with open(image_path, 'rb') as f:
                image_data = f.read()
            
            # 图片数据转Base64编码
            image_base64 = base64.b64encode(image_data).decode('utf-8')
            
            # 获取百度OCR API访问令牌
            access_token = await OCRClient._get_baidu_access_token(
                settings.BAIDU_API_KEY, 
                settings.BAIDU_SECRET_KEY
            )
            
            if "error" in access_token:
                return {"error": f"获取百度访问令牌失败: {access_token['error']}"}
            
            # 调用通用文字识别高精度版API
            ocr_url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/accurate?access_token={access_token['token']}"
            
            async with ClientSession() as session:
                headers = {
                    'Content-Type': 'application/x-www-form-urlencoded'
                }
                
                params = {
                    'image': image_base64,
                    'language_type': 'CHN_ENG',  # 中英文混合
                    'detect_direction': 'true',  # 检测图片朝向
                    'probability': 'true'  # 返回识别结果中每行的置信度
                }
                
                async with session.post(ocr_url, headers=headers, data=params) as response:
                    response_json = await response.json()
                    
                    if "error_code" in response_json:
                        return {"error": f"百度OCR API调用失败: {response_json.get('error_msg', '未知错误')}"}
                    
                    return response_json
        
        except Exception as e:
            logging.error(f"百度OCR API调用异常: {str(e)}")
            return {"error": f"OCR处理异常: {str(e)}"}
    
    @staticmethod
    async def _get_baidu_access_token(api_key: str, secret_key: str) -> Dict:
        """
        获取百度AI开放平台访问令牌
        
        Args:
            api_key: 百度API Key
            secret_key: 百度Secret Key
            
        Returns:
            Dict: 包含访问令牌的字典，或错误信息
        """
        try:
            from aiohttp import ClientSession
            
            # 令牌获取URL
            token_url = f"https://aip.baidubce.com/oauth/2.0/token"
            
            async with ClientSession() as session:
                params = {
                    'grant_type': 'client_credentials',
                    'client_id': api_key,
                    'client_secret': secret_key
                }
                
                async with session.post(token_url, params=params) as response:
                    response_json = await response.json()
                    
                    if "error" in response_json:
                        return {"error": f"获取令牌失败: {response_json.get('error_description', '未知错误')}"}
                    
                    return {"token": response_json.get("access_token")}
        
        except Exception as e:
            logging.error(f"获取百度访问令牌异常: {str(e)}")
            return {"error": f"获取访问令牌异常: {str(e)}"}
    
    @staticmethod
    async def extract_text(ocr_result: Dict) -> str:
        """
        从OCR结果中提取文本内容
        
        Args:
            ocr_result: OCR识别结果
            
        Returns:
            str: 提取的文本内容
        """
        try:
            text_lines = []
            
            if "words_result" in ocr_result:
                for word_info in ocr_result["words_result"]:
                    if "words" in word_info:
                        text_lines.append(word_info["words"])
            
            if not text_lines:
                return "未能从图片中提取有效文本内容"
            
            # 组合成文本
            combined_text = "\n".join(text_lines)
            # 清理文本
            cleaned_text = OCRClient._clean_text(combined_text)
            
            return cleaned_text
        
        except Exception as e:
            logging.error(f"提取OCR文本异常: {str(e)}")
            return f"提取文本异常: {str(e)}"
    
    @staticmethod
    async def process_images_with_concurrency(
        image_paths: List[str],
        max_concurrent: int = settings.MAX_CONCURRENT_OCR
    ) -> Dict[str, str]:
        """
        使用有限的并发数处理多张图片
        """
        results = {}
        semaphore = asyncio.Semaphore(max_concurrent)
        
        async def process_image(img_path: str):
            async with semaphore:
                ocr_result = await OCRClient.recognize_image(img_path)
                text = await OCRClient.extract_text(ocr_result)
                return img_path, text
        
        # 创建所有任务
        tasks = [process_image(img_path) for img_path in image_paths]
        
        # 等待所有任务完成
        for task in asyncio.as_completed(tasks):
            img_path, text = await task
            results[img_path] = text
        
        return results
    
    @staticmethod
    def _clean_text(text: str) -> str:
        """
        清理OCR提取的文本
        
        Args:
            text: 原始OCR文本
            
        Returns:
            str: 清理后的文本
        """
        # 去除多余的空格
        cleaned = re.sub(r'\s+', ' ', text)
        # 去除特殊字符，保留基本标点
        cleaned = re.sub(r'[^\w\s,.，。!?！？:：;；()\[\]【】\-—]', '', cleaned)
        # 合并多个换行
        cleaned = re.sub(r'\n+', '\n', cleaned)
        
        return cleaned.strip()
    
    @staticmethod
    async def process_image(image_path: str) -> str:
        """
        处理图片，提取文本内容
        
        Args:
            image_path: 图片路径
            
        Returns:
            str: 提取的文本内容
        """
        try:
            print(f"使用Tesseract OCR处理图片: {image_path}")
            # 首先尝试使用本地Tesseract
            text = await OCRClient._use_tesseract(image_path)
            if text:
                return text
        except Exception as e:
            print(f"Tesseract OCR不可用: {str(e)}")
        
        try:
            print(f"尝试使用百度OCR API处理图片: {image_path}")
            # 如果本地OCR失败，使用百度OCR API
            text = await OCRClient._use_baidu_ocr(image_path)
            if text:
                return text
        except Exception as e:
            print(f"百度OCR API调用异常: {str(e)}")
            print(f"百度OCR API处理失败: {str(e)}")
        
        # 如果所有OCR方法都失败，生成模拟文本
        mock_text = OCRClient._generate_mock_text()
        print(f"生成模拟OCR文本，长度: {len(mock_text)}")
        print(f"所有OCR方法均失败，使用模拟文本: {mock_text[:100]}")
        return mock_text
    
    @staticmethod
    async def _use_tesseract(image_path: str) -> str:
        """使用Tesseract OCR"""
        try:
            import pytesseract
            from PIL import Image
            
            # 打开图片
            image = Image.open(image_path)
            # 使用Tesseract进行OCR
            text = pytesseract.image_to_string(image, lang='chi_sim+eng')
            return text.strip()
        except Exception as e:
            raise Exception(f"Tesseract OCR处理失败: {str(e)}")
    
    @staticmethod
    async def _use_baidu_ocr(image_path: str) -> str:
        """使用百度OCR API"""
        try:
            from aip import AipOcr
            
            # 初始化百度OCR客户端
            client = AipOcr(
                settings.BAIDU_APP_ID,
                settings.BAIDU_API_KEY,
                settings.BAIDU_SECRET_KEY
            )
            
            # 读取图片
            with open(image_path, 'rb') as fp:
                image = fp.read()
            
            # 使用asyncio运行同步方法    
            loop = asyncio.get_running_loop()
            result = await loop.run_in_executor(None, lambda: client.basicGeneral(image))
            
            if 'words_result' in result:
                # 提取文本
                texts = [item['words'] for item in result['words_result']]
                return '\n'.join(texts)
            else:
                raise Exception("OCR结果中没有文本内容")
                
        except Exception as e:
            raise Exception(f"OCR处理异常: {str(e)}")
    
    @staticmethod
    def _generate_mock_text() -> str:
        """生成模拟文本（当所有OCR方法都失败时使用）"""
        return """品牌: 根据图片推测的品牌名称
该图片可能展示了产品的外观、包装或详细信息。
由于OCR处理失败，无法提取具体文本内容。
建议重新上传更清晰的图片或手动输入产品信息。"""
    
    @staticmethod
    async def process_images_with_concurrency(image_paths: List[str], max_concurrency: int = 5) -> Dict[str, str]:
        """
        并发处理多个图片，提取文本内容
        
        Args:
            image_paths: 图片路径列表
            max_concurrency: 最大并发数
            
        Returns:
            Dict[str, str]: 图片路径和对应的文本内容的字典
        """
        # 创建信号量控制并发
        semaphore = asyncio.Semaphore(max_concurrency)
        
        async def process_with_semaphore(img_path):
            async with semaphore:
                text = await OCRClient.process_image(img_path)
                return img_path, text
        
        # 创建所有任务
        tasks = [process_with_semaphore(img_path) for img_path in image_paths]
        
        # 等待所有任务完成
        results = await asyncio.gather(*tasks)
        
        # 将结果转换为字典
        return {os.path.basename(img_path): text for img_path, text in results}
    
    @staticmethod
    async def process_image_with_api(image_path: str, api_url: Optional[str] = None) -> str:
        """
        使用API处理图片OCR（备用方法，可配置外部OCR服务）
        
        Args:
            image_path: 图片路径
            api_url: API地址，如果为None则使用本地处理
            
        Returns:
            str: 提取的文本内容
        """
        if not api_url:
            return await OCRClient.process_image(image_path)
        
        try:
            # 读取图片并转为base64
            with open(image_path, "rb") as img_file:
                base64_image = base64.b64encode(img_file.read()).decode('utf-8')
            
            # 发送API请求
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    api_url,
                    json={"image": base64_image}
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        if result.get("success") and "text" in result:
                            return OCRClient._clean_text(result["text"])
                        else:
                            return f"API处理失败: {result.get('message', '未知错误')}"
                    else:
                        return f"API请求失败，状态码: {response.status}"
        except Exception as e:
            logging.error(f"OCR API处理失败: {str(e)}")
            # 如果API失败，回退到本地处理
            return await OCRClient.process_image(image_path)
    
    @staticmethod
    async def batch_process_product_images(product_folder_path: str) -> Dict[str, str]:
        """
        批量处理产品文件夹下的所有图片
        
        Args:
            product_folder_path: 产品文件夹路径
            
        Returns:
            Dict[str, str]: 图片文件名和OCR文本的映射
        """
        # 获取文件夹下所有图片
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        image_paths = []
        
        try:
            for file_name in os.listdir(product_folder_path):
                file_path = os.path.join(product_folder_path, file_name)
                if os.path.isfile(file_path) and any(file_name.lower().endswith(ext) for ext in image_extensions):
                    image_paths.append(file_path)
            
            logging.info(f"找到 {len(image_paths)} 张产品图片待处理")
            
            if not image_paths:
                return {"error": "未找到有效图片文件"}
            
            # 并发处理所有图片
            results = await OCRClient.process_images_with_concurrency(
                image_paths,
                max_concurrency=settings.MAX_CONCURRENT_OCR
            )
            
            # 将结果转换为 {文件名: 文本内容} 格式
            return {os.path.basename(path): text for path, text in results.items()}
            
        except Exception as e:
            logging.error(f"批量处理产品图片失败: {str(e)}")
            return {"error": f"批量处理产品图片失败: {str(e)}"}
    
    @staticmethod
    async def save_ocr_results(product_folder_path: str, ocr_results: Dict[str, str]) -> Tuple[bool, str]:
        """
        将OCR结果保存到产品分析结果文件
        
        Args:
            product_folder_path: 产品文件夹路径
            ocr_results: OCR结果，格式为 {图片文件名: OCR文本}
            
        Returns:
            Tuple[bool, str]: (是否成功, 消息)
        """
        try:
            result_file = os.path.join(product_folder_path, "analysis_result.yaml")
            
            # 如果文件已存在，先读取现有内容
            existing_data = {}
            if os.path.exists(result_file):
                with open(result_file, 'r', encoding='utf-8') as f:
                    existing_data = yaml.safe_load(f) or {}
            
            # 更新或添加OCR结果
            if "image_texts" not in existing_data:
                existing_data["image_texts"] = {}
            
            existing_data["image_texts"].update(ocr_results)
            
            # 保存更新后的内容
            with open(result_file, 'w', encoding='utf-8') as f:
                yaml.dump(existing_data, f, allow_unicode=True, default_flow_style=False, Dumper=yaml.SafeDumper)
            
            return True, f"成功保存 {len(ocr_results)} 张图片的OCR结果"
            
        except Exception as e:
            logging.error(f"保存OCR结果失败: {str(e)}")
            return False, f"保存OCR结果失败: {str(e)}" 